import os
import json
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit


def assign_region(lat, lon):
    """
    根据纬度粗略划分区域：陕北 / 关中 / 陕南
    注意：这里只是示例，实际可以细化调整。
    """
    if lat > 37:
        return "陕北"
    elif 33 <= lat <= 37:
        return "关中"
    else:
        return "陕南"


# 自定义角度转换函数
def dms_to_deg(dms_str):
    """
    将形如 '1074322' 的字符串转为十进制度（度分秒格式：DDDMMSS）
    """
    dms = int(dms_str)
    deg = dms // 10000
    minutes = (dms % 10000) // 100
    seconds = dms % 100
    return deg + minutes / 60 + seconds / 3600

# 降水分组
precipitation_bins = [0.1, 50, 100, 150, np.inf]
bin_labels = ["0.1_50mm", "50_100mm", "100_150mm", "150mm+"]

# 读取数据
def read_data_files(data_dir):
    all_records = []
    for file_name in os.listdir(data_dir):
        if file_name.endswith(".json"):
            file_path = os.path.join(data_dir, file_name)
            with open(file_path, "r", encoding="utf-8") as f:
                data = json.load(f)
            for item in data:
                try:
                    lon = dms_to_deg(item["Lon"])
                    lat = dms_to_deg(item["Lat"])
                    pre = float(item["PRE_1h"])
                    all_records.append({
                        "lon": lon,
                        "lat": lat,
                        "pre": pre
                    })
                except Exception:
                    continue
    return pd.DataFrame(all_records)

# 计算空间点之间的距离（欧氏距离）
def compute_distances(x, y):
    n = len(x)
    distances = np.zeros((n, n))
    for i in range(n):
        for j in range(n):
            distances[i, j] = np.sqrt((x[i] - x[j])**2 + (y[i] - y[j])**2)
    return distances

# 计算半变异函数
def compute_variogram(x, y, z, bins=10):
    # 计算所有点之间的距离
    distances = compute_distances(x, y)
    
    # 定义半变异函数的 bin
    lags = np.linspace(0, np.max(distances), bins)
    semivariance = np.zeros(bins-1)
    
    # 计算每个 bin 内的半变异值
    for i in range(bins-1):
        # 找到距离在当前 bin 区间的点对
        mask = (distances >= lags[i]) & (distances < lags[i+1])
        # 计算半变异函数
        semivariance[i] = np.nanmean((z[:, None] - z[None, :]) ** 2 * mask)
    
    return lags[:-1], semivariance

# 球形模型
def spherical_model(h, nugget, sill, range_):
    return nugget + sill * (1.5 * (h / range_) - 0.5 * (h / range_)**3) * (h < range_)

# 使用 scipy.curve_fit 来拟合半变异函数
from scipy.optimize import curve_fit

# 三种常见模型定义
def spherical_model(h, nugget, sill, range_):
    h = np.asarray(h)
    gamma = np.where(
        h <= range_,
        nugget + (sill - nugget) * (1.5 * h / range_ - 0.5 * (h / range_) ** 3),
        sill
    )
    return gamma

def exponential_model(h, nugget, sill, range_):
    return nugget + (sill - nugget) * (1 - np.exp(-h / range_))

def gaussian_model(h, nugget, sill, range_):
    return nugget + (sill - nugget) * (1 - np.exp(-(h / range_) ** 2))

model_functions = {
    "spherical": spherical_model,
    "exponential": exponential_model,
    "gaussian": gaussian_model
}

# 半变异函数拟合函数
def fit_variogram(x, y, z, model_list=["spherical", "exponential", "gaussian"]):
    from scipy.spatial.distance import pdist, squareform

    coords = np.vstack((x, y)).T
    dists = pdist(coords)
    values = pdist(z.reshape(-1, 1), lambda u, v: (u - v) ** 2 / 2)
    
    # 分 bin 计算实验半变异函数
    n_lags = 15
    max_dist = np.percentile(dists, 30)  # 排除极远点
    bins = np.linspace(0, max_dist, n_lags + 1)
    bin_centers = 0.5 * (bins[:-1] + bins[1:])
    semivariance = []

    for i in range(n_lags):
        mask = (dists >= bins[i]) & (dists < bins[i + 1])
        if np.any(mask):
            semivariance.append(np.mean(values[mask]))
        else:
            semivariance.append(np.nan)

    semivariance = np.array(semivariance)

    # 清理 NaN
    valid = ~np.isnan(semivariance)
    lags = bin_centers[valid]
    semivariance = semivariance[valid]

    best_model = None
    best_params = None
    best_rmse = np.inf

    for model_name in model_list:
        model_func = model_functions[model_name]

        try:
            # 拟合曲线，添加参数约束
            params, _ = curve_fit(
                model_func,
                lags,
                semivariance,
                p0=[0.1, np.max(semivariance), np.max(lags) / 2],
                bounds=([0, 0, 1e-3], [np.inf, np.inf, np.inf])
            )

            fitted = model_func(lags, *params)
            rmse = np.sqrt(np.mean((semivariance - fitted) ** 2))

            if rmse < best_rmse:
                best_model = model_name
                best_params = params
                best_rmse = rmse

        except Exception as e:
            print(f"[警告] 模型 {model_name} 拟合失败: {e}")
            continue

    if best_model is None:
        print("[错误] 所有模型拟合失败，跳过该分组")
        return None, None, None, None

    return best_model, best_params, (lags, semivariance), None

# 绘制拟合图
def plot_fit(experimental_variogram, model_params, bin_label):
    lags = experimental_variogram[0]
    experimental = experimental_variogram[1]
    theoretical = spherical_model(lags, *model_params)

    plt.figure(figsize=(10, 6))
    plt.scatter(lags, experimental, label="Experimental Variogram")
    plt.plot(lags, theoretical, label=f"Fitted {model_name} Model", color='red')
    plt.xlabel("Lag Distance")
    plt.ylabel("Semivariance")
    plt.title(f"Variogram Fit: {bin_label}")
    plt.legend()
    plt.grid(True)
    os.makedirs("nihe", exist_ok=True)
    save_path = os.path.join("nihe", f"{bin_label}_variogram_fit.png")
    plt.savefig(save_path)
    plt.close()

# 主程序
def main():
    data_dir = "data"  # JSON 文件目录
    df = read_data_files(data_dir)

    df = df[(df["pre"] > 0.1) & (df["pre"] < 500)]   # 过滤无效数据
    df["precipitation_group"] = pd.cut(df["pre"], bins=precipitation_bins, labels=bin_labels)
    df["region"] = df.apply(lambda row: assign_region(row["lat"], row["lon"]), axis=1)

    os.makedirs("nihe", exist_ok=True)

    # 遍历每个区域
    for region in df["region"].unique():
        region_df = df[df["region"] == region]
        print(f"\n[区域] {region} 数据点总数: {len(region_df)}")

        all_params = {}

        for bin_label in bin_labels:
            group_df = region_df[region_df["precipitation_group"] == bin_label]
            print(f"[处理] 区域 {region} 分组 {bin_label} 样本数: {len(group_df)}")

            if len(group_df) < 10:
                print(f"[跳过] 区域 {region} 分组 {bin_label} 样本过少")
                continue

            x = group_df["lon"].values
            y = group_df["lat"].values
            z = group_df["pre"].values

            model_name, params, experimental, _ = fit_variogram(x, y, z)

            if params is None:
                continue

            nugget, sill, range_ = params
            print(f"拟合参数：nugget={nugget:.4f}, sill={sill:.4f}, range={range_:.4f}")

            plot_fit(experimental, params, f"{region}_{bin_label}")

            all_params[bin_label] = {
                "model": model_name,
                "nugget": float(nugget),
                "sill": float(sill),
                "range": float(range_)
            }

        # 保存该区域的拟合参数
        save_path = os.path.join("nihe", f"variogram_params_{region}.json")
        with open(save_path, "w", encoding="utf-8") as f:
            json.dump(all_params, f, ensure_ascii=False, indent=4)
        print(f"[完成] 区域 {region} 拟合参数保存至 {save_path}")
   
if __name__ == "__main__":
    main()
